
import os
from options.test_options import TestOptions
from data import create_dataset
from models import create_model
from util.visualizer import save_images
from util import save_html
from tqdm import tqdm
import util.util as util
import torch
from thop import profile
from thop import clever_format
from util.ssim import cal_ssim_bt_dirs
from tqdm import tqdm

from pathlib import Path

from models.cut_model import CUTModel


if __name__ == '__main__':
    opt = TestOptions().parse()  # get test options
    # hard-code some parameters for test
    opt.num_threads = 0   # test code only supports num_threads = 1
    opt.batch_size = 1    # test code only supports batch_size = 1
    # disable data shuffling; comment this line if results on randomly chosen images are needed.
    opt.serial_batches = True
    # no flip; comment this line if results on flipped images are needed.
    opt.no_flip = True
    # no visdom display; the test code saves the results to a HTML file.
    opt.display_id = -1
    # create a dataset given opt.dataset_mode and other options
    dataset = None
    # train_dataset = None
    if opt.phase == "test":
        dataset = create_dataset(opt)
    else:
        dataset = create_dataset(util.copyconf(opt, phase="train"))
    test_flag = True


    ##############
    # g = CUTModel(opt)
    # input = torch.randn(1, 3, 768, 1280).cuda()
    # macs, params = profile(g, inputs=(input, ))  # macs = 0.5 * flops
    # macs, params = clever_format([macs, params], "%.3f")
    # print("macs, params:", macs, params)
    ##############
    start_epoch = opt.epoch_count
    total_epoch = 50
    for epoch in range(start_epoch, total_epoch+1, 10):
        opt.epoch_count = epoch
        opt.epoch = epoch
        # create a webpage for viewing the results
        web_dir = os.path.join(opt.results_dir, opt.name, '{}_{}'.format(
        opt.phase, opt.epoch_count))  # define the website directory
        print('creating web directory = %s, Epoch = %s'%( web_dir, opt.epoch_count))
        webpage = save_html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (
        opt.name, opt.phase, opt.epoch_count))
        for i, data in tqdm(enumerate(dataset)):
            if i == 0:
                    # create a model given opt.model and other options
                model = create_model(opt)
                model.data_dependent_initialize(data)
                # regular setup: load and print networks; create schedulers
                model.setup(opt)
                model.parallelize()
                if opt.eval:
                    model.eval()
            model.set_input(data)  # unpack data from data loader
            model.test()           # run inference
            visuals = model.get_test_visuals()  # get image results
            img_path = model.get_image_paths()     # get image paths
            save_images(webpage, visuals, img_path, width=opt.display_winsize)
        # cal ssim and psnr  
        fake_B_dir = Path(web_dir) / 'images/fake_B'
        real_A_dir = Path(opt.dataroot) / 'test/A/'
        if opt.phase == "test":
            ssim, psnr = cal_ssim_bt_dirs(fake_B_dir, real_A_dir)
        else:
            data_B_dir = Path(opt.dataroot) / 'train/B'
            fid_command = "CUDA_VISIBLE_DEVICES=" + str(opt.gpu_ids) + " python -m pytorch_fid " + str(fake_B_dir) + " " + str(data_B_dir)
            print("calc FID bewteen", fake_B_dir, data_B_dir)
            os.system(fid_command)
    print("------------------------------------------------------")
